# !/usr/bin/python3
# -*- coding:utf-8 -*-
# Copyright 2022 The Chinaunicom Software Team. All rights reserved.
# @Author : dyu
# @Time   : 2022-8-12

### 模型推理模块 ###

import os
import time
import warnings

warnings.simplefilter(action='ignore', category=FutureWarning)
import tensorflow as tf
import pandas as pd
from tensorflow.python.keras.backend import set_session
from src.intelligent_interaction.engine.annoy_model import AnnoyRecall
from src.intelligent_interaction.engine.sem_predictor import SimbertEmbeding
from src.intelligent_interaction.engine.sem_conf import *


# memoryList = list(map(int, os.popen("nvidia-smi -q -d Memory | grep -A4 GPU | grep Total | awk '{print $3}'").readlines()))
# memoryTotal = memoryList[0]
# memoryLimited = 2000.0  # use 3G memory


class BasePredictor(object):
    def __init__(self, task_name):
        # log_server.logging('===============Prediction Process Beginning !===============')
        self.task_name = task_name
        tf_config = tf.ConfigProto()
        #         tf_config.gpu_options.per_process_gpu_memory_fraction = memoryLimited / memoryTotal
        self.sess = tf.Session(config=tf_config)
        self.graph = tf.get_default_graph()
        set_session(self.sess)
        self.simbert, self.annoy_0013, self.annoy_t8, self.annoy_t9, self.annoy_stand = self.create_model()
        print('模型加载完毕')

    def create_model(self):
        """
        创建模型
        """
        simbert = SimbertEmbeding()
        annoy_0013 = AnnoyRecall(f_dim=config['f_dim'], ann_path=config['pro_id']['faq_0013']['ann_path'],
                                 source=config['pro_id']['faq_0013']['source_path'])
        annoy_t8 = AnnoyRecall(f_dim=config['f_dim'], ann_path=config['pro_id']['faq_t8']['ann_path'],
                               source=config['pro_id']['faq_t8']['source_path'])
        annoy_t9 = AnnoyRecall(f_dim=config['f_dim'], ann_path=config['pro_id']['faq_t9']['ann_path'],
                               source=config['pro_id']['faq_t9']['source_path'])
        annoy_stand = AnnoyRecall(f_dim=config['f_dim'], ann_path=config['pro_id']['faq_stand']['ann_path'],
                                  source=config['pro_id']['faq_stand']['source_path'])
        return simbert, annoy_0013, annoy_t8, annoy_t9, annoy_stand

    def predict(self, question, pro_id='stand'):
        """
        预测过程
        :paramter question:  输入问题
        :paramter pro_id  :  省份编码，default faq_stand
        """
        try:
            ### 第一步：bert embedding ###
            emb_vec = self.simbert.single_embedding(question)
        except Exception as e:
            return []

        ### 第二步：Annoy召回 ###
        print("pro_id:", pro_id)
        if pro_id == '0013':
            print('---0013---')
            predict = self.annoy_0013.vecter_recall_predict(vector=emb_vec, n_sims=config['n_sims'],
                                                            search_k=config['search_k'],
                                                            confidence=config['pro_id']['faq_0013']['confidence'])
        elif pro_id == 't8':
            print('---t8---')
            predict = self.annoy_t8.vecter_recall_predict(vector=emb_vec, n_sims=config['n_sims'],
                                                          search_k=config['search_k'],
                                                          confidence=config['pro_id']['faq_t8']['confidence'])
        elif pro_id == 't9':
            print('---t9---')
            predict = self.annoy_t9.vecter_recall_predict(vector=emb_vec, n_sims=config['n_sims'],
                                                          search_k=config['search_k'],
                                                          confidence=config['pro_id']['faq_t9']['confidence'])
        elif pro_id == 'stand':
            print('---stand---')
            predict = self.annoy_stand.vecter_recall_predict(vector=emb_vec, n_sims=config['n_sims'],
                                                             search_k=config['search_k'],
                                                             confidence=config['pro_id']['faq_stand']['confidence'])
        else:
            print('---other---')
            predict = []
        return predict


if __name__ == '__main__':
    faq = BasePredictor('faq')
    result = faq.predict(question='罚金是啥，怎么搞', pro_id='0013')
    print(result)
